import os
import argparse
from simnibs import sim_struct, run_simnibs

def run_leadfield_simulation(subpath, leadfield_name=None):
    """
    Runs a SimNIBS leadfield simulation for tDCS.

    Parameters:
    - subpath: str, required
        Path to the subject's folder.
    - leadfield_name: str, required
        Name for the leadfield. This will be used to set the output directory.
    """
    if leadfield_name==None:
        leadfield_name=f"leadfield"
    # Initialize the TDCSLEADFIELD object
    tdcs_lf = sim_struct.TDCSLEADFIELD()
    
    # Set the subject path and output path
    tdcs_lf.subpath = subpath
    tdcs_lf.pathfem = os.path.join(subpath, leadfield_name)
    os.makedirs(tdcs_lf.pathfem,exist_ok = True)
    # Set solver options
    #tdcs_lf.solver_options = 'pardiso'
    
    # Run the simulation
    run_simnibs(tdcs_lf)
    print(f"Leadfield simulation completed. Outputs saved to {tdcs_lf.pathfem}")


    leadfield_name+="_volume"
    # Initialize the TDCSLEADFIELD object
    tdcs_lf = sim_struct.TDCSLEADFIELD()
    # Set the subject path and output path
    tdcs_lf.subpath = subpath
    tdcs_lf.pathfem = os.path.join(subpath, leadfield_name)
    os.makedirs(tdcs_lf.pathfem, exist_ok=True)
    tdcs_lf.tissues= [1,2,6]
    tdcs_lf.interpolation =None
    
    # Set solver options
    #tdcs_lf.solver_options = 'pardiso'
    
    # Run the simulation
    run_simnibs(tdcs_lf)
    print(f"Leadfield simulation completed. Outputs saved to {tdcs_lf.pathfem}")

def main():
    # Set up argument parsing
    parser = argparse.ArgumentParser(description="Run a SimNIBS tDCS leadfield simulation.")
    
    parser.add_argument("--subpath", type=str, required=True, help="Path to the subject's folder.")
    parser.add_argument("--leadfield_name", type=str,help="Name for the leadfield (used to set the output directory).(Optional)")
    
    # Parse the arguments
    args = parser.parse_args()
    
    # Run the leadfield simulation with the provided arguments
    run_leadfield_simulation(args.subpath, args.leadfield_name)

if __name__ == "__main__":
    main()
